import torch
from torch.utils.tensorboard import SummaryWriter  # 画出训练损失图像
from train_dataset import train_dataset
from cnn import Net

# 检测是否有可用的GPU，如果没有则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载训练数据，实例化SummaryWriter，将日志文件存放到log文件夹，该文件夹会自动创建在项目目录
data_load = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True)
writer = SummaryWriter('log')

# 实例化网络、设置损失函数和优化器
model = Net().to(device)  # 将模型发送到指定设备
model.train()
loss_function = torch.nn.CrossEntropyLoss()
optim = torch.optim.SGD(model.parameters(), lr=1e-3)

for epoch in range(20):
    step = 0
    losses = []
    for img, target in data_load:
        img, target = img.to(device), target.to(device)  # 将数据发送到指定设备
        pred = model(img)
        loss = loss_function(pred, target.long())  # 这里要注意target要使用长整型。
        optim.zero_grad()
        loss.backward()
        optim.step()

        step += 1
        print(f'epoch:{epoch + 1} [{step * 8}/{len(data_load) * 8}] loss:{loss / 8}')
        losses.append(loss)
    mean_loss = sum(losses) / step * 8
    writer.add_scalar('training loss', mean_loss, epoch)  # 每个epoch计算平均损失，然后画图

# 保存模型
torch.save(model.state_dict(), 'Net.pth')